package problems.practice;

import java.util.*;

/**
 * 787. K 站中转内最便宜的航班
 * https://leetcode.cn/problems/cheapest-flights-within-k-stops/
 *
 * @author habitplus
 * @since 2022/8/14 21:03
 */
public class T787 {
    // BFS + 剪枝  ----> TME
    public int findCheapestPrice1(int n, int[][] flights, int src, int dst, int k) {
        if (n < 1 || flights == null || src == dst || k < 0) return -1;
        HashMap<Integer, List<int[]>> map = new HashMap<>();
        for (int[] f : flights) {
            List<int[]> list = map.getOrDefault(f[0], null);
            if (list == null) {
                list = new ArrayList<>();
                list.add(new int[]{f[1], f[2]});
                map.put(f[0], list);
            } else {
                list.add(new int[] {f[1], f[2]});
            }


        }

        Queue<int[]> que = new LinkedList<>();
        que.add(new int[]{src, 0});
        int i, dis;
        int ret = Integer.MAX_VALUE;
        while (!que.isEmpty() && k + 1 > 0) {
            for (i = que.size(); i > 0; --i) {
                int[] s = que.poll();

                List<int[]> list = map.getOrDefault(s[0], null);
                if (list != null) {
                    for (int[] e : list) {
                        dis = s[1] + e[1];
                        // 剪枝一：如果发现比之前的距离短，才有可能进行中转
                        if (dis < ret) {
                            if (e[0] == dst) ret = dis;
                            else {
                                // 剪枝二：不是终点才中转
                                que.offer(new int[] {e[0], dis});
                            }
                        }
                    }
                }
            }
            k--;
        }

        return ret == Integer.MAX_VALUE ? -1 : ret;
    }

    int INF = 1000007;

    public int findCheapestPrice(int n, int[][] flights, int src, int dst, int k) {
        return bfs(n, flights, src, dst, k);
    }

    private int bfs(int n, int[][] flights, int src, int dst, int k) {
        // 整理题目给定的flights，转换成每个节点的子节点有哪些
        List<int[]>[] g = new List[n];
        for (int i = 0; i < n; i++)
            g[i] = new ArrayList<>();


        for (int[] flight : flights) {
            g[flight[0]].add(new int[]{flight[1], flight[2]});
        }

        // 表示src到i到最小价格
        int[] ans = new int[n];
        Arrays.fill(ans, INF);
        Queue<int[]> queue = new LinkedList<>();
        queue.offer(new int[]{src, 0});
        // 退出条件加上 k 的限制
        while (!queue.isEmpty() && k + 1 > 0) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                int[] poll = queue.poll();
                for (int[] path : g[poll[0]]) {
                    int distance = poll[1] + path[1];
                    // 剪枝1，小于 i 之前记录的最小值，且小于 dst 之前记录的最小值
                    if (distance < ans[path[0]] && distance < ans[dst]) {
                        ans[path[0]] = distance;
                        // 剪枝2，到 dst 了就不用继续往下了
                        if (path[0] != dst) {
                            queue.offer(new int[]{path[0], distance});
                        }
                    }
                }
            }
            k--;
        }
        return ans[dst] >= INF ? -1 : ans[dst];
    }

}
